import os
import re
import json
import argparse
from typing import List, Any

import requests


# 可配置变量：文章地址（默认指向本地articles目录中的文件）
ARTICLE_PATH = None  # 如需默认文章路径，可设置为例如："articles/sample.md"


# ------------------------------
# 基础文件读取工具
# ------------------------------

SUPPORTED_EXTS = {".md", ".txt", ".html", ".htm"}


def read_text_file(path: str, max_chars: int = 16000) -> str:
    encodings = ["utf-8", "utf-8-sig", "gb18030", "latin-1"]
    for enc in encodings:
        try:
            with open(path, "r", encoding=enc, errors="ignore") as f:
                content = f.read()
                return content[:max_chars]
        except Exception:
            continue
    return ""


def html_to_text(s: str) -> str:
    try:
        import html as _html
    except Exception:
        _html = None
    s = re.sub(r"<script[\s\S]*?</script>", " ", s, flags=re.IGNORECASE)
    s = re.sub(r"<style[\s\S]*?</style>", " ", s, flags=re.IGNORECASE)
    s = re.sub(r"<[^>]+>", " ", s)
    s = re.sub(r"\s+", " ", s).strip()
    if _html:
        s = _html.unescape(s)
    return s


def read_article_content(article_path: str) -> str:
    if not article_path or not os.path.isfile(article_path):
        raise FileNotFoundError(f"找不到文章文件: {article_path}")
    _, ext = os.path.splitext(article_path.lower())
    text = read_text_file(article_path, max_chars=16000)
    if ext in {".html", ".htm"}:
        return html_to_text(text)
    return text


# ------------------------------
# AI 调用（与 use_gemini_search.py 同风格）
# ------------------------------

def call_ai_chat(prompt: str, api_key: str, model: str = "gpt-4o-mini", base_url: str = "https://aihubmix.com/v1/chat/completions") -> Any:
    import os as _os
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    data = {
        "model": model,
        "messages": [
            {"role": "user", "content": prompt}
        ]
    }
    def _should_use_local_proxy() -> bool:
        keys = (
            "USE_PROXY", "FORCE_LOCAL_PROXY", "LOCAL_PROXY",
            "HTTP_PROXY", "HTTPS_PROXY", "ALL_PROXY",
            "http_proxy", "https_proxy", "all_proxy",
        )
        return any(bool(_os.environ.get(k)) for k in keys)
    def _local_proxies():
        addr = (_os.environ.get("LOCAL_PROXY") or "127.0.0.1:7897").strip()
        return {"http": f"http://{addr}", "https": f"http://{addr}"}
    s = requests.Session()
    s.trust_env = False
    if _should_use_local_proxy():
        s.proxies = _local_proxies()
    resp = s.post(base_url, headers=headers, data=json.dumps(data), timeout=60)
    resp.raise_for_status()
    result = resp.json()
    content = result["choices"][0]["message"]["content"]

    if isinstance(content, str):
        s = content.strip()
        if s.startswith("```json") and s.endswith("```"):
            s = s[7:-3].strip()
        elif s.startswith("```") and s.endswith("```"):
            s = s[3:-3].strip()
        try:
            parsed = json.loads(s)
            return parsed
        except Exception:
            return content
    return content


# ------------------------------
# 提示词构建
# ------------------------------

def build_prompt(article_text: str) -> str:
    prompt = f"""
你是插图描述生成助手。请基于 “文章内容” 判断需要配图的主题，从文章核心主题关联的不同场景、使用情境、细节视角、关联物品等维度拓展画面，确保图片种类丰富且与文章内容贴合不突兀，最终输出用于图像生成的自然语言描述列表。

严格要求：

输出必须是 “JSON 数组（字符串列表）”，每个元素对应一张图的描述；不要输出任何额外文字或说明；不要包裹代码块。
每张图的人物数量仅能为 1 或 2（若无需人物，请用背景 / 场景描述，避免出现任何人）。
每条描述需简洁明确，可直接用于图片生成，包含：主体（可结合主题关联的不同核心元素，而非单一物体）、场景 / 背景（如主题相关的展示场景、制作场景、使用场景等）、关键物体、风格 / 光线 / 色调、镜头角度等信息；禁止出现 “文字水印、Logo、拼贴、过多文字”。
避免敏感、暴力、成人、侵权等元素。
所有描述需围绕文章核心主题展开，同时确保不同描述的画面主体、情境或视角不重复，体现多样性，且融入文章时无突兀感。
生成的描述要丰富不要停留在单一物体或场景，展开联想。

文章内容：\n"""
    prompt += article_text.strip()[:12000]
    prompt += "\n\n仅返回JSON数组（字符串列表），不要生成其他任何内容。"
    return prompt


# ------------------------------
# 主流程函数
# ------------------------------

def generate_image_descriptions(article_path: str,
                                api_key: str,
                                model: str = "gpt-4o-mini") -> List[str]:
    text = read_article_content(article_path)
    prompt = build_prompt(text)
    result = call_ai_chat(prompt, api_key=api_key, model=model)

    # 解析与校验：需要返回列表[str]
    if isinstance(result, list):
        cleaned: List[str] = []
        for item in result:
            if isinstance(item, str):
                s = item.strip()
                if s:
                    cleaned.append(s)
        return cleaned

    # 如果不是列表，尝试从字符串中解析JSON
    if isinstance(result, str):
        try:
            parsed = json.loads(result)
            if isinstance(parsed, list):
                return [str(x).strip() for x in parsed if isinstance(x, (str, int, float))]
        except Exception:
            pass

    # 兜底：返回空列表
    return []


def generate_descriptions(article_path):
    parser = argparse.ArgumentParser(description="生成文章插图描述（列表）")
    parser.add_argument("--article-path", type=str, default=ARTICLE_PATH, help="文章文件路径（.md/.txt/.html）")
    parser.add_argument("--api-key", type=str, default=None, help="AI API Key（默认读取 AIHUBMIX_API_KEY 环境变量）")
    parser.add_argument("--model", type=str, default="gpt-4o-mini", help="模型名称，与 use_gemini_search.py 风格一致")
    parser.add_argument("--out", type=str, default=None, help="可选：将列表以JSON格式保存到文件")

    args = parser.parse_args()

    api_key = "sk-PjfrbJ5yZ1kAeSH6875c4fD5551f4f4b8e839d28Da91A6Db"
    if not api_key:
        raise SystemExit("缺少 API Key。请通过 --api-key 或环境变量 AIHUBMIX_API_KEY 提供。")
    if not article_path:
        raise SystemExit("请通过 --article-path 指定文章文件路径，或在脚本顶部设置 ARTICLE_PATH。")

    descriptions = generate_image_descriptions(article_path=article_path, api_key=api_key, model=args.model)
    return descriptions

    # if args.out:
    #     with open(args.out, "w", encoding="utf-8") as f:
    #         json.dump(descriptions, f, ensure_ascii=False, indent=2)
    #     print(f"已保存到 {args.out}")
    # else:
    #     # 直接打印纯JSON数组
    #     print(json.dumps(descriptions, ensure_ascii=False))


if __name__ == "__main__":
    article_path = "yellow_diamond_guide_revised.md"
    generate_descriptions(article_path)

